import torch
from utils import HARDatasetPt, HARDatasetMat
from mymodel import TransformerSequenceClassifier
import lightning as L
import lightning.pytorch.loggers
import lightning.pytorch.accelerators
import os
from torchinfo import summary

train_ds = HARDatasetPt("./combined_train_dataset.pt")
train_dl = torch.utils.data.DataLoader(train_ds, batch_size=64, shuffle=True)

test_ds = HARDatasetMat("./UCIHAR_Test.mat")
test_dl = torch.utils.data.DataLoader(test_ds, batch_size=64, shuffle=False)


class TransformerSequenceClassifierL(L.LightningModule):
    def __init__(self):
        super().__init__()
        self.model = TransformerSequenceClassifier()
        self.loss = torch.nn.CrossEntropyLoss()

    def forward(self, x, *args, **kwargs):
        return self.model(x)

    def training_step(self, batch, batch_idx, *args, **kwargs):
        x, y = batch
        y_hat = self.model(x)
        loss = self.loss(y_hat, y)
        self.log('train_loss', loss, prog_bar=True)
        return loss

    def validation_step(self, batch, batch_idx, *args, **kwargs):
        x, y = batch
        y_hat = self.model(x)
        loss = self.loss(y_hat, y)
        self.log('val_loss', loss, prog_bar=True)
        return loss

    def configure_optimizers(self):
        optimizer = torch.optim.Adam(self.model.parameters(), lr=1e-5)
        return optimizer


model = TransformerSequenceClassifierL()
root_path = "lightning_logs"

summary(model.model, input_size=(1, 128, 6))

checkpoint_callback = L.pytorch.callbacks.ModelCheckpoint(
    dirpath=os.path.join(root_path, "checkpoints"),
    filename="{epoch:02d}-{val_loss:.2f}",
    monitor="val_loss",
    verbose=True,
    save_top_k=3,
    mode="min",
    save_weights_only=True
)

early_stopping_callback = L.pytorch.callbacks.EarlyStopping(
    monitor="val_loss",
    patience=10,
    verbose=True,
    mode="min"
)

rich_progress_bar = L.pytorch.callbacks.RichProgressBar(
    refresh_rate=1,
    leave=True
)

logger = lightning.pytorch.loggers.TensorBoardLogger(
    save_dir=os.path.join(root_path, "logs"),
    name="mymodel",
    log_graph=False
)



trainer = L.Trainer(
    accelerator="auto",
    precision=32,
    max_epochs=20,
    callbacks=[checkpoint_callback, early_stopping_callback, rich_progress_bar],
    logger=logger,
    devices=1
)

trainer.fit(
    model=model,
    train_dataloaders=train_dl,
    val_dataloaders=test_dl
)
